import { useCallback, useRef, useSyncExternalStore } from 'react';
import type { TypedKey } from '../context/types';
import { globalContext } from '../globalContext';
import { shallowEqual } from '../utils/shallowEqual';

export function useListen<T>(
  key: TypedKey<T>,
  equalityFn: (a: T, b: T) => boolean = shallowEqual,
): [T, (value: T | ((prev: T) => T)) => void] {
  const set = useCallback(
    (newValue: T | ((prev: T) => T)) => {
      const finalValue =
        typeof newValue === 'function'
          ? (newValue as (prev: T) => T)(globalContext.getData(key))
          : newValue;

      globalContext.setData(key, finalValue);
    },
    [key],
  );

  // 使用React 18的useSyncExternalStore优化性能
  const snapshot = useSyncExternalStore(
    // 订阅函数
    (onStoreChange) => {
      const unsubscribe = globalContext.subscribe(key, (_newValue) => {
        onStoreChange();
      });
      return unsubscribe;
    },
    // 获取当前值
    () => globalContext.getData(key),
    // 获取服务器端渲染的初始值
    () => globalContext.getData(key),
  );

  // 使用ref保存前一个值，只有当值真正变化时才返回新值
  const stableSnapshot = useRef<T>(snapshot);

  // 如果值有变化，更新稳定快照
  if (!equalityFn(stableSnapshot.current, snapshot)) {
    stableSnapshot.current = snapshot;
  }

  return [stableSnapshot.current, set];
}
